from utils import criterions
from utils.tensorbackend import get_backend
import numpy
import matplotlib.pyplot as plt


class ComparatorBase:
    def __init__(self, algos: tuple):
        """

        :param algos:
        """
        for algo in algos:
            if not isinstance(algo, tuple):
                raise TypeError("The algos parameter should be a tuple containing a (algorithm class:Subclass of "
                                "AlgorithmBase, name:string) tuple")
        self.algos = algos

    def compareAndPlot(self, noiseGenerator, n, L, repeat=1, criterion=criterions.L2Norm, input=None, weight_ori=None,
                       weight_target=None, title="PlainComparatorOutput", needLog: bool = True, backend=None,
                       device="cpu"):
        """

        :param noiseGenerator:
        :param n:
        :param L:
        :param repeat:
        :param criterion:
        :param input:
        :param weight_ori:
        :param weight_target:
        :param backend:
        :param device:
        :param title:
        :param needLog:
        :return:
        """
        raise NotImplementedError("The method in an abstract class is not invokable")


class EvenComparator(ComparatorBase):
    """
    This comparator is to compare all the algorithms in the completely same input, target, noise and initial weights
    """

    def compareAndPlot(self, noiseGenerator, n, L, repeat=1, criterion=criterions.L2Norm, input=None, weight_ori=None,
                       weight_target=None, title="PlainComparatorOutput", needLog: bool = True, backend=None,
                       device="cpu"):
        if repeat < 1:
            raise ValueError(f"repeat time cannot be less than 1,{repeat} got instead")

        res = self._compareAndPlot(noiseGenerator, n, L, criterion, input, weight_ori, weight_target, backend, device)
        for i in range(2, repeat):
            tmp = self._compareAndPlot(noiseGenerator, n, L, criterion, input, weight_ori, weight_target, backend,
                                       device)
            for key in tmp:
                res[key] = [res[key][i] + tmp[key][i] for i in range(L)]

        for key in res:
            for i in range(L):
                res[key][i] = res[key][i] / L

        x = numpy.linspace(0, L, L).tolist()
        fig, ax = plt.subplots()
        for algo in self.algos:
            if needLog:
                ax.plot(x, numpy.log10(res[algo[1]]) * 20, label=algo[1])
            else:
                ax.plot(x, res[algo[1]], label=algo[1])
        ax.set_xlabel("iterations")
        ax.set_ylabel("error(20lgx)" if needLog else "error")
        ax.set_title(title)
        ax.legend(loc='upper right')
        plt.show()

    def _compareAndPlot(self, noiseGenerator, n, L, criterion, input=None, weight_ori=None, weight_target=None,
                        backend=None, device="cpu"):
        if backend is None:
            if input is not None:
                backend = get_backend(input)
            elif weight_ori is not None:
                backend = get_backend(weight_ori)
            elif weight_target is not None:
                backend = get_backend(weight_target)
            else:
                backend = get_backend(numpy.zeros(1))
        result_dict = {}
        if input is None:
            input = backend.normal(loc=0, scale=1, shape=(n, L), device=device)
        if weight_ori is None:
            weight_ori = backend.normal(loc=0, scale=1, shape=(n, 1), device=device)
        if weight_target is None:
            weight_target = backend.normal(loc=0, scale=1, shape=(n, 1), device=device)

        noise = noiseGenerator.getNoise(shape=(1, L), backend=backend, device=device)
        target = backend.matmul(weight_target.T, input) + noise
        for algo in self.algos:
            result_dict[algo[1]] = []
            algo[0].weight = weight_ori
            for i in range(0, L):
                algo[0].iterate(target=backend.unsqueeze(target[:, i], dim=1),
                                input=backend.unsqueeze(input[:, i], dim=1))
                result_dict[algo[1]].append(backend.to_numpy(criterion(algo[0].weight, weight_target)))

        return result_dict
